package org.rbio.paper.methylglm.dmr;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.TreeMap;

import org.kohsuke.args4j.CmdLineParser;
import org.kohsuke.args4j.Option;

// class to build DMRs
public class DMRClusterBuilder {

	// program arguments
	@Option(name = "-dev_file", usage = "CpG methylation file with deviance values", required = true)
	private String dev_file;

	@Option(name = "-min_dev", usage = "minimum deviance", required = true)
	private double min_dev;

	@Option(name = "-win", usage = "window size", required = true)
	private int win;

	@Option(name = "-out", usage = "output DMR file", required = true)
	private String out;

	// make the clusters from the deviance values
	public static void main(String args[]) throws Exception {
		DMRClusterBuilder builder = new DMRClusterBuilder();

		// parse command line
		CmdLineParser parser = new CmdLineParser(builder);
		try {
			parser.parseArgument(args);
		} catch (Exception e) {
			System.err.println(e.getMessage());
			System.err
					.println("Make differentially methylated region [DMR] from the calculated deviance files");
			parser.printUsage(System.err);
			return;
		}
		builder.process();
	}

	// process [ref pos ... dev]
	public void process() throws Exception {
		Map<String, PriorityQueue<CpG>> qu_tab = new TreeMap<String, PriorityQueue<CpG>>();

		// make the priority queue
		List<CpG> cpg_lst = new ArrayList<CpG>();
		BufferedReader ibuf = new BufferedReader(new FileReader(dev_file));
		String line = ibuf.readLine();
		while ((line = ibuf.readLine()) != null) {
			// empty line
			if (line.length() == 0) {
				continue;
			}

			// parse line
			String arr[] = line.split("\t");
			String ref = arr[0];
			int pos = Integer.parseInt(arr[1]);
			double dev = Double.parseDouble(arr[arr.length - 1]);

			// check for minimum deviance
			if (dev >= min_dev) {
				// fill the queue
				PriorityQueue<CpG> qu = qu_tab.get(ref);
				if (qu == null) {
					qu = new PriorityQueue<CpG>();
					qu_tab.put(ref, qu);
				}
				CpG cpg = new CpG(pos, dev);
				qu.add(cpg);
				cpg_lst.add(cpg);
			}
		}
		CpGComparator pos_compare = new CpGComparator();
		Collections.sort(cpg_lst, pos_compare);
		ibuf.close();

		// make cluster
		BufferedWriter obuf = new BufferedWriter(new FileWriter(out));
		obuf.write("ref\tstart\tend\tcnt\tpos\n");
		for (String ref : qu_tab.keySet()) {
			PriorityQueue<CpG> qu = qu_tab.get(ref);
			while (!qu.isEmpty()) {
				// poll the CpG
				CpG cpg = qu.poll();

				// make the cluster
				List<CpG> clust = new ArrayList<CpG>();
				clust.add(cpg);
				int idx = cpg_lst.indexOf(cpg);

				// look in the right direction
				int start = cpg.pos;
				for (int i = idx + 1; i < cpg_lst.size(); i++) {
					CpG rt_cpg = cpg_lst.get(i);
					if (rt_cpg.pos - start <= win) {
						// CpG within win bp of anchor
						clust.add(rt_cpg);
						qu.remove(rt_cpg);
						start = rt_cpg.pos;
					} else {
						// CpG is not within win bp of the anchor
						break;
					}
				}

				// look in the left direction
				start = cpg.pos;
				for (int i = idx - 1; i >= 0; i--) {
					CpG lt_cpg = cpg_lst.get(i);
					if (start - lt_cpg.pos <= win) {
						// CpG within win bp of anchor
						clust.add(lt_cpg);
						qu.remove(lt_cpg);
						start = lt_cpg.pos;
					} else {
						// CpG is not within win bp of the anchor
						break;
					}
				}

				// write the cluster
				Collections.sort(clust, pos_compare);
				write(ref, clust, obuf);
			}
		}
		obuf.close();
	}

	// write the cluster
	public void write(String ref, List<CpG> clust, BufferedWriter obuf)
			throws Exception {
		CpG first_cpg = clust.get(0);
		CpG last_cpg = clust.get(clust.size() - 1);
		String pos_lst = Integer.toString(first_cpg.pos);
		for (int i = 1; i < clust.size(); i++) {
			pos_lst += "," + clust.get(i).pos;
		}
		obuf.write(ref + "\t" + first_cpg.pos + "\t" + last_cpg.pos + "\t"
				+ clust.size() + "\t" + pos_lst + "\n");
	}

	// CpG class
	public static class CpG implements Comparable<CpG> {
		int pos;
		double deviance;

		// constructor
		public CpG(int pos, double deviance) {
			this.pos = pos;
			this.deviance = deviance;
		}

		// make it comparable by deviance
		public int compareTo(CpG that) {
			int res = 0;
			if (this.deviance > that.deviance) {
				res = 1;
			} else if (this.deviance < that.deviance) {
				res = -1;
			}
			return res * -1;
		}
	}

	// comparator class
	public static class CpGComparator implements Comparator<CpG> {

		// compare two CpGs by position
		public int compare(CpG cpg1, CpG cpg2) {
			int res = 0;
			if (cpg1.pos > cpg2.pos) {
				res = 1;
			} else if (cpg1.pos < cpg2.pos) {
				res = -1;
			}
			return res;
		}
	}
}
